'''
Description: 计算iou
Author: suyunzheng
Date: 2021-11-25 19:14:27
LastEditTime: 2021-11-25 22:33:43
LastEditors: maple
'''
import enum
from operator import ipow
from sklearn.metrics import confusion_matrix


from os.path import exists, join
import numpy as np
import time
from sklearn.metrics import confusion_matrix


import numpy as np 

def iou(y_pred, y_true):
    y_pred = y_pred.flatten()
    y_true = y_true.flatten()

    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2])
    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2,3,4,5,6,7,8,9,10,11,12])
    # current = confusion_matrix(y_true, y_pred, labels=[0,1,2,3,4,5,6,7,8,9,10,11,12])
    current = confusion_matrix(y_true, y_pred, labels=[0,1,2])

    intersection =  np.diag(current)
    gt = current.sum(axis = 1)
    pred = current.sum(axis = 0)
    union = gt + pred - intersection
    ioU = intersection / union.astype(np.float32)+1e-8
    return np.mean(ioU)



def log_out(out_str, f_out):
    f_out.write(out_str + '\n')
    f_out.flush()
    print(out_str)

class IoUCalculator:
    def __init__(self, class_num):
        self.class_num = class_num
        self.gt_classes = [0 for _ in range(self.class_num)]
        self.positive_classes = [0 for _ in range(self.class_num)]
        self.true_positive_classes = [0 for _ in range(self.class_num)]

    def add_data(self, pred_labels, gt_labels):
        if not isinstance(pred_labels, list):
            pred_labels = np.argmax(pred_labels, axis=1)
        # val_total_correct = 0
        # val_total_seen = 0

        # correct = np.sum(pred_labels == gt_labels)
        # val_total_correct += correct
        # val_total_seen += len(gt_labels)
        conf_matrix = confusion_matrix(gt_labels, pred_labels, labels=np.arange(0, self.class_num, 1))
        self.gt_classes += np.sum(conf_matrix, axis=1)
        self.positive_classes += np.sum(conf_matrix, axis=0)
        self.true_positive_classes += np.diagonal(conf_matrix)

    def compute_iou(self):
        iou_list = []
        for n in range(0, self.class_num, 1):
            if float(self.gt_classes[n] + self.positive_classes[n] - self.true_positive_classes[n]) != 0:
                iou = self.true_positive_classes[n] / float(self.gt_classes[n] + 
                        self.positive_classes[n] - self.true_positive_classes[n])
                iou_list.append(iou)
            else:
                iou_list.append(0.0)
        mean_iou = sum(iou_list) / float(self.class_num)
        return mean_iou, iou_list

    def compute_acc(self):
        iou_list = []
        for n in range(0, self.class_num, 1):
            if float( self.positive_classes[n] ) != 0:
                iou = self.true_positive_classes[n] / float( self.positive_classes[n] )
                iou_list.append(iou)
            else:
                iou_list.append(0.0)
        mean_iou = sum(iou_list) / float(self.class_num)
        return mean_iou, iou_list

    def print_iou(self, CLASS_NAME):
        miou, iou_list = self.compute_iou()
        print("===> miou:{}".format(miou))
        print("===> iou <===")
        for i, ele in enumerate(zip(CLASS_NAME, iou_list)):
            # print("{}\t\t\t{}".format(ele[0], ele[1]))
            print('%s : %.2f%%' % (ele[0], ele[1]*100))

        print("===> iou <===")
        
if __name__== "__main__":
    print("numpy")

    a = [1,0,2,0,1,2]
    pred = [[0.9, 0.05, 0.05],
            [0.9, 0.05, 0.05],
            [0.9, 0.05, 0.05],
            [0.9, 0.05, 0.05],
            [0.9, 0.05, 0.05],
            [0.9, 0.05, 0.05]]
    pred = np.asarray(pred).reshape(-1,3)
    # a  = np.array(a)
    # pred = np.array(pred)

    # print("a.shape:{}".format(a.shape))
    # print("pred shape:{}".format(pred.shape))
    # iou_value = iou(pred, a)
    # print("iou:{}".format(iou_value))


    iou_cla = IoUCalculator(3)
    iou_cla.add_data(pred, a)
    re = iou_cla.compute_iou()
    print("re:{}".format(re))
